import matplotlib.pyplot as plt
import math
import os
import numpy as np

# Define a function to read data
def read_data(file_path):
    data = []
    x_coords = []
    y_coords = []
    with open(file_path, 'r') as file:
        for line in file:
            if line.strip() == "":
                if x_coords and y_coords:
                    data.append((x_coords, y_coords))
                    x_coords = []
                    y_coords = []
            else:
                try:
                    x, y = map(float, line.split())
                    x_coords.append(x)
                    y_coords.append(y)
                except ValueError:
                    # Ignore lines that do not contain two float values
                    continue
    if x_coords and y_coords:
        data.append((x_coords, y_coords))
    
    return data

# Define a function to plot scatter points and lines
def plot_data(x_coords, y_coords, label, color):
    plt.scatter(x_coords, y_coords, color=color, label=label, s=5) 
    plt.plot(x_coords, y_coords, color=color)  # Add lines

# Get all .txt files in the output directory that do not end with 'x'
output_dir = 'output'

A_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('A_') and f.endswith('.txt') and not f.endswith('x.txt')]
Ae_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('Ae') and f.endswith('.txt') and not f.endswith('x.txt')]
C_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('C') and f.endswith('.txt') and not f.endswith('x.txt')]
D_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('D') and f.endswith('.txt') and not f.endswith('x.txt')]
E_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('E') and f.endswith('.txt') and not f.endswith('x.txt')]
F_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('F') and f.endswith('.txt') and not f.endswith('x.txt')]
T4_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('T4') and f.endswith('.txt') and not f.endswith('x.txt')]
T5_paths = [f for f in os.listdir(output_dir) 
              if f.startswith('T5') and f.endswith('.txt') and not f.endswith('x.txt')]

# Define colors for plotting
colors = ['g', 'r', 'b', 'y', 'c', 'm', 'k']
    
# Plot each A_ file
Name_A=['N=6','N=11','N=21','N=41','N=81']
for i, A_path in enumerate(A_paths):
    fig, ax = plt.subplots()
    data = read_data(os.path.join(output_dir, A_path))
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, Name_A[j % len(Name_A)], colors[j % len(colors)])
    # Add the standard curve 1/(1+25x^2)
    x_standard = np.linspace(-1.1, 1.1, 400)
    y_standard = 1 / (1 + 25 * x_standard**2)
    ax.plot(x_standard, y_standard, color='k', label='1/(1+25x^2)')
    ax.legend()
    ax.set_title(f"{os.path.splitext(A_path)[0]}")
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(os.path.join('images', f'{os.path.splitext(A_path)[0]}.png'))
    plt.close(fig)

Name_Ae = ['Linear', 'Cubic']
# Plot each Ae file
for i, Ae_path in enumerate(Ae_paths):
    data = read_data(os.path.join(output_dir, Ae_path))
    fig, ax = plt.subplots()
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, Name_Ae[j % len(Name_Ae)], colors[j % len(colors)])
    ax.legend()
    ax.set_title("ln(Errors) vs ln(N)")
    ax.set_xlabel('ln(N)')
    ax.set_ylabel('ln(Errors)')
    plt.savefig(os.path.join('images', 'A_Convergences.png'))
    plt.close(fig)

# Plot each C file
Name_C=['Linear Bspline','Complete cubic Bspline']
for i, C_path in enumerate(C_paths):
    fig, ax = plt.subplots()
    data = read_data(os.path.join(output_dir, C_path))
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, Name_C[j % len(Name_C)], colors[j % len(colors)])
    # Add the standard curve 1/(1+25x^2)
    x_standard = np.linspace(-5.1, 5.1, 400)
    y_standard = 1 / (1 + x_standard**2)
    ax.plot(x_standard, y_standard, color='k', label='1/(1+x^2)')
    ax.legend()
    ax.set_title(f"{os.path.splitext(C_path)[0]}")
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(os.path.join('images', f'{os.path.splitext(C_path)[0]}.png'))
    plt.close(fig)
    
# Plot each D file
Name_D=['Linear Bspline','Domplete cubic Bspline']
for i, D_path in enumerate(D_paths):
    fig, ax = plt.subplots()
    data = read_data(os.path.join(output_dir, D_path))
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, Name_D[j % len(Name_D)], colors[j % len(colors)])
    # Add the standard curve 1/(1+25x^2)
    ax.legend()
    ax.set_title(f"{os.path.splitext(D_path)[0]}")
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(os.path.join('images', f'{os.path.splitext(D_path)[0]}.png'))
    plt.close(fig)

# Plot each E file
Name_E = ['Bezier', 'CubicPP']
for i, E_path in enumerate(E_paths):
    data = read_data(os.path.join(output_dir, E_path))
    fig, ax = plt.subplots()
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, Name_E[j % len(Name_E)], colors[j % len(colors)])
    
    # Add the standard curve based on the filename
    if E_path[3] == '1':
        t_standard = np.linspace(0, 2 * np.pi, 500)
        x_standard = np.sqrt(3) * np.cos(t_standard)
        y_standard = (2.0 / 3.0) * np.sqrt(np.abs(np.sqrt(3) * np.cos(t_standard))) + (2.0 / 3.0) * np.sqrt(3) * np.sin(t_standard)
        ax.plot(x_standard, y_standard, color='k', label='R1 Standard Curve')
    elif E_path[3] == '2':
        t_standard = np.linspace(0, 6 * np.pi, 500)
        x_standard = np.sin(t_standard) + t_standard * np.cos(t_standard)
        y_standard = np.cos(t_standard) - t_standard * np.sin(t_standard)
        ax.plot(x_standard, y_standard, color='k', label='R2 Standard Curve')
    
    ax.legend()
    ax.set_title(f"{os.path.splitext(E_path)[0]}")
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(os.path.join('images', f'{os.path.splitext(E_path)[0]}.png'))
    plt.close(fig)

# Plot each F file
for i, F_path in enumerate(F_paths):
    data = read_data(os.path.join(output_dir, F_path))
    num_datasets = len(data)
    num_rows = math.ceil(math.sqrt(num_datasets))
    fig, axs = plt.subplots(num_rows, num_rows, figsize=(15, num_rows * 5))
    j = 0 
    for col in range(num_rows):
        for row in range(num_rows):
            ax = axs[row, col]
            if row >= col:
                x_coords, y_coords = data[j]
                ax.plot(x_coords, y_coords, color=colors[(j) % len(colors)], label=f"Part {j+1}")
                
                ax.legend()
                ax.set_title(f"{os.path.splitext(F_path)[0]} - Part {j + 1}")
                ax.set_xlabel('X-axis')
                ax.set_ylabel('Y-axis')
                j = j + 1
            else:
                ax.axis('off')
    # Save the plot with the same name as the .txt file but with .png extension
    fig.suptitle(f"n={i+1} Divided difference table")
    plt.savefig(os.path.join('images', f'{os.path.splitext(F_path)[0]}.png'))
    plt.close(fig)

Name_T4 = ['PPForm', 'Bspline']
# Plot each T4 file
for i, T4_path in enumerate(T4_paths):
    data = read_data(os.path.join(output_dir, T4_path))
    fig, ax = plt.subplots()
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, Name_T4[j % len(Name_T4)], colors[j % len(colors)])
    ax.legend()
    ax.set_title(f"{os.path.splitext(T4_path)[0]}")
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(os.path.join('images', f'{os.path.splitext(T4_path)[0]}.png'))
    plt.close(fig)

# Plot each T5 file
for i, T5_path in enumerate(T5_paths):
    data = read_data(os.path.join(output_dir, T5_path))
    fig, ax = plt.subplots()
    for j, (x_coords, y_coords) in enumerate(data):
        plot_data(x_coords, y_coords, os.path.splitext(T5_path)[0], colors[j % len(colors)])
    ax.legend()
    ax.set_title(f"{os.path.splitext(T5_path)[0]}")
    ax.set_xlabel('X-axis')
    ax.set_ylabel('Y-axis')
    plt.savefig(os.path.join('images', f'{os.path.splitext(T5_path)[0]}.png'))
    plt.close(fig)